import mindspore.nn as nn
import mindspore as ms
from mindvision.dataset import Mnist
from mindvision.classification.models import lenet
from mindvision.engine.callback import LossMonitor
# 处理MNIST数据集
mnist = Mnist("./mnist", split="train", batch_size=32, resize=32)
dataset_train = mnist.run()
network = lenet(num_classes=10, pretrained=False)
# 定义损失函数
net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
# 定义优化器函数
net_opt = nn.Momentum(network.trainable_params(), learning_rate=0.01, momentum=0.9)
# 设置模型保存参数，模型训练保存参数的step为1875。
config_ck = ms.CheckpointConfig(save_checkpoint_steps=1875, keep_checkpoint_max=10)
# 应用模型保存参数
ckpoint = ms.ModelCheckpoint(prefix="lenet", directory="./lenet", config=config_ck)
# 初始化模型参数
model = ms.Model(network, loss_fn=net_loss, optimizer=net_opt, metrics={'accuracy'})
# 训练网络模型，并保存为lenet-1_1875.ckpt文件
model.train(1, dataset_train, callbacks=[ckpoint, LossMonitor(0.01, 1875)])
